import pandas as pd
import re
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.optimizers import Adam
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error


# 读取数据集
file_path = r"./data/beijing.csv"
data = pd.read_csv(file_path)

# 选择需要的特征和目标变量
features = [
    "constructionTime",
    "square",
    "livingRoom",
    "drawingRoom",
    "kitchen",
    "bathRoom",
    "floor",
    "subway",
]
target = "price"

# 提取相关列
data = data[features + [target]]

# 数据清洗
# 去除缺失值
data.dropna(inplace=True)
# 去除包含字符串“未知”和提取floor的数值部分，去除空值的记录
data = data[
    (data["constructionTime"] != "未知")
    & (data["bathRoom"].notnull())
    & (data["livingRoom"].notnull())
]
data["floor"] = data["floor"].apply(
    lambda x: re.findall(r"\d+", str(x))[0] if re.findall(r"\d+", str(x)) else None
)

# 转换数据类型
data["constructionTime"] = pd.to_datetime(
    data["constructionTime"], errors="coerce"
).dt.year
data["square"] = pd.to_numeric(data["square"], errors="coerce")
data["livingRoom"] = pd.to_numeric(data["livingRoom"], errors="coerce")
data["drawingRoom"] = pd.to_numeric(data["drawingRoom"], errors="coerce")
data["kitchen"] = pd.to_numeric(data["kitchen"], errors="coerce")
data["bathRoom"] = pd.to_numeric(data["bathRoom"], errors="coerce")
data["floor"] = pd.to_numeric(data["floor"], errors="coerce")
data["subway"] = pd.to_numeric(data["subway"], errors="coerce")
data["price"] = pd.to_numeric(data["price"], errors="coerce")

# 再次去除因转换错误导致的缺失值
data.dropna(inplace=True)
print(data.info())
print(data.head(10))


# 分离特征和目标变量
X = data[features]
y = data[target]

# 分割数据集为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# 标准化特征
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# 构建MLP模型
model = Sequential(
    [
        Dense(128, input_dim=X_train.shape[1], activation="relu"),
        Dropout(0.2),
        Dense(64, activation="relu"),
        Dropout(0.2),
        Dense(32, activation="relu"),
        Dropout(0.2),
        Dense(1),
    ]
)

# 编译模型
model.compile(optimizer=Adam(learning_rate=0.001), loss="mean_squared_error")

# 打印模型摘要
model.summary()

# 设置回调函数
early_stopping = tf.keras.callbacks.EarlyStopping(
    monitor="val_loss", patience=10, restore_best_weights=True
)
reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
    monitor="val_loss", factor=0.2, patience=5, min_lr=0.001
)

# 训练模型
history = model.fit(
    X_train,
    y_train,
    epochs=100,
    validation_split=0.2,
    callbacks=[early_stopping, reduce_lr],
)

# 评估模型
loss = model.evaluate(X_test, y_test)
print(f"\nTest Loss: {loss}")

# 预测并计算评估指标
y_pred = model.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
rmse = np.sqrt(mse)
mae = mean_absolute_error(y_test, y_pred)
print(f"Test MSE: {mse}")
print(f"Test RMSE: {rmse}")
print(f"Test MAE: {mae}")

# 保存模型
model.save("price_model.h5")
